import os
import json
import time
from flask import Flask, request, jsonify
from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageContentType, ChatEventType, COZE_CN_BASE_URL

app = Flask(__name__)
# Set JSON encoder to ensure proper Chinese character handling
app.json.ensure_ascii = False

# Initialize Coze client
coze = Coze(auth=TokenAuth("pat_n1G8J1PfymCiBSi2yrSh6afMxk6LK1qrXl9D7JZPDh2cTG6LoD2n4SrHJ95xNAvv"),base_url=COZE_CN_BASE_URL)

# Create a bot instance in Coze, copy the last number from the web link as the bot's ID.
BOT_ID = "7445582723185557504"
USER_ID = "test123"

# Rate limiter configuration
RATE_LIMIT_SECONDS = 5
last_request_time = 0

def ask_coze(user_question_text: str):
    # Rate limiting check
    global last_request_time
    current_time = time.time()
    time_since_last_request = current_time - last_request_time
    
    if time_since_last_request < RATE_LIMIT_SECONDS:
        return jsonify({
            'error': '请求过于频繁，请等待 {} 秒后重试'.format(
                round(RATE_LIMIT_SECONDS - time_since_last_request, 1)
            )
        }), 429
    
    last_request_time = current_time

    try:

        # Create chat and get response
        chat_poll = coze.chat.create_and_poll(
            bot_id=BOT_ID,
            user_id=USER_ID,
            additional_messages=[
                Message.build_user_question_text(user_question_text),
            ],
        )

        # Get the response message
        response_content = None
        for msg in chat_poll.messages:
            if msg.content:
                response_content = msg.content

        # Parse the JSON response
        ans = json.loads(chat_poll.messages[1].content)

        # Include token usage in response if chat is completed
        token_usage = None
        if chat_poll.chat.status == ChatStatus.COMPLETED:
            token_usage = chat_poll.chat.usage.token_count

        return {
            'result': ans,
            'token_usage': token_usage
        }

    except json.JSONDecodeError:
        return {'error': 'JSON 格式错误'}
    except Exception as e:
        return {'error': str(e)}

if __name__ == "__main__":
    print(ask_coze("番茄炒蛋"))